{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e80da62",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from onekey_algo.custom.components.metrics import calc_dice, calc_iou\n",
    "\n",
    "def seg_eval(pred, label, clss=[0, 1]):\n",
    "    \"\"\"\n",
    "    calculate the dice between prediction and ground truth\n",
    "    input:\n",
    "        pred: predicted mask\n",
    "        label: groud truth\n",
    "        clss: eg. [0, 1] for binary class\n",
    "    \"\"\"\n",
    "    Ncls = len(clss)\n",
    "    eval_matric = [None] * Ncls\n",
    "    [depth, height, width] = pred.shape\n",
    "    for idx, cls in enumerate(clss):\n",
    "        # binary map\n",
    "        pred_cls = np.zeros([depth, height, width], dtype=np.uint8)\n",
    "        pred_cls[np.where(pred == cls)] = 1\n",
    "        label_cls = np.zeros([depth, height, width], dtype=np.uint8)\n",
    "        label_cls[np.where(label == cls)] = 1\n",
    "\n",
    "        metric = [calc_dice(pred_cls, label_cls), calc_iou(pred_cls, label_cls)]\n",
    "        eval_matric[idx] = metric\n",
    "\n",
    "    return eval_matric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff69afba",
   "metadata": {},
   "outputs": [],
   "source": [
    "from glob import glob\n",
    "import SimpleITK as sitk\n",
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from onekey_algo import get_param_in_cwd\n",
    "\n",
    "models = os.listdir(get_param_in_cwd('save_dir'))\n",
    "cohort_metric = []\n",
    "metric_spec = []\n",
    "tn = None\n",
    "for model in models:\n",
    "    root = os.path.join(get_param_in_cwd('save_dir'), model, f'infer')\n",
    "    metric_names = ['Dice', 'mIOU']\n",
    "    all_metrics = []\n",
    "    samples = []\n",
    "    for fs in glob(os.path.join(get_param_in_cwd('data_root'), 'masks', '*.nii.gz')):\n",
    "        gt_mask = fs\n",
    "        pred_mask = os.path.join(root, os.path.basename(gt_mask))\n",
    "        if not os.path.exists(pred_mask):\n",
    "#             print(f'{pred_mask}不存在！')\n",
    "            continue\n",
    "        samples.append(os.path.basename(gt_mask))\n",
    "        all_metrics.append(seg_eval(sitk.GetArrayFromImage(sitk.ReadImage(pred_mask)),\n",
    "                                    sitk.GetArrayFromImage(sitk.ReadImage(gt_mask)), clss=[0, 1]))\n",
    "    metric = pd.DataFrame(np.mean(np.array(all_metrics), axis=1), columns=metric_names)\n",
    "    cohort_metric.append(pd.DataFrame(metric.mean(axis=0)).T)\n",
    "    metric['model'] = model\n",
    "    info = pd.concat([pd.DataFrame(samples, columns=['ID']), metric], axis=1)\n",
    "    metric_spec.append(info)\n",
    "info = pd.concat(metric_spec, axis=0)\n",
    "os.makedirs('results', exist_ok=True)\n",
    "info.to_csv(f'results/metrics_details.csv', index=False)\n",
    "info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59930479",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = info.groupby(['model']).agg('mean').reset_index()\n",
    "metrics.to_csv('results/metrics.csv', index=False)\n",
    "metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "389e2401",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
